import torch
from transformers import AutoTokenizer
from bigdl.llm.transformers import AutoModel
from bigdl.llm.optimize import load_low_bit, low_memory_init
from flask import Flask
from flask_socketio import SocketIO

app = Flask(__name__)
socketio = SocketIO(app, async_mode='threading', cors_allowed_origins="*", logger=True)

with low_memory_init():
    model = AutoModel.from_pretrained("chatglm2-6b", trust_remote_code=True)

model_path = "chatglm2-6b_lowbit"
model = load_low_bit(model, model_path)
tokenizer = AutoTokenizer.from_pretrained("chatglm2-6b", trust_remote_code=True)

CHATGLM_V2_PROMPT_TEMPLATE = "问：{prompt}\n\n答："

conversation_history = {}
Identity = "你的身份是一个课程助教，你叫VidBot"

@socketio.on('message')

def handle_message(data):
    
    msg=data.get("message")
    student_id=data.get("student_id")
    print(f"用户 {student_id} :{msg}")  # 打印用户输入

    user_chat_his=[]
    if student_id in conversation_history:
        user_chat_his=conversation_history[f"{student_id}"]
    else:
        user_chat_his.append(Identity.strip())
        conversation_history[f"{student_id}"]=[]

    # 将用户输入添加到聊天历史
    user_chat_his.append(msg)

    model_response = generate_response(user_chat_his)
    print(conversation_history)
    socketio.emit('message', {'data': f'机器人：{model_response}'})

def generate_response(user_chat_his):
    user_chat_his = '\n'.join([CHATGLM_V2_PROMPT_TEMPLATE.format(prompt=utterance) for utterance in user_chat_his])
    input_ids = tokenizer.encode(user_chat_his, return_tensors="pt")

    with torch.inference_mode():
        output = model.generate(input_ids, max_length=1024)
        model_response = tokenizer.decode(output[0], skip_special_tokens=True)

    model_response = model_response[len(user_chat_his):].strip()  # 删除生成结果中的 prompt

    return model_response

@socketio.on('disconnect')
def disconnect(sid):
    
    # student_id=data.get("student_id")
    # if student_id in conversation_history:
    #     del conversation_history[f"{student_id}"]

    # return 'OK'
    print(sid)

if __name__ == '__main__':
    # 使用 SSL 启动 SocketIO 服务器
    socketio.run(app,host='0.0.0.0', port=8765, ssl_context=('/www/wwwroot/default/certificate.crt', '/www/wwwroot/default/private.key'))
